#!/bin/bash
#
# Klpmake is a linux kernel livepatch making tool.
#
# Copyright (c) 2023 laokz <zhangkai@iscas.ac.cn>
# Klpmake is licensed under Mulan PSL v2.
# You can use this software according to the terms and conditions of
# the Mulan PSL v2. You may obtain a copy of Mulan PSL v2 at:
#          http://license.coscl.org.cn/MulanPSL2
# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES
# OF ANY KIND, EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED
# TO NON-INFRINGEMENT, MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
# See the Mulan PSL v2 for more details.

set -e
LOCALE=C
KLPMAKE_DIR=$(dirname $0)
cwd=$PWD
klpsyms=_klpmake.syms
klpconf=_klpsrc.conf
srcroot=
dbgroot=
debug=
declare -a syscalls

case $(uname -m) in
    aarch64)    SYSCALL_PREFIX="__arm64_sys_"
                SYSCALL_ARG0="regs[0]"
                SYSCALL_ARG1="regs[1]"
                SYSCALL_ARG2="regs[2]"
                SYSCALL_ARG3="regs[3]"
                SYSCALL_ARG4="regs[4]"
                SYSCALL_ARG5="regs[5]" ;;
    riscv64)    SYSCALL_PREFIX="__riscv_sys_"
                SYSCALL_ARG0="orig_a0"
                SYSCALL_ARG1="a1"
                SYSCALL_ARG2="a2"
                SYSCALL_ARG3="a3"
                SYSCALL_ARG4="a4"
                SYSCALL_ARG5="a5" ;;
    x86_64)     SYSCALL_PREFIX="__x64_sys_"
                SYSCALL_ARG0="di"
                SYSCALL_ARG1="si"
                SYSCALL_ARG2="dx"
                SYSCALL_ARG3="r10"
                SYSCALL_ARG4="r8"
                SYSCALL_ARG5="r9" ;;
    *)          echo "ERROR: unsupported arch"; return 6 ;;
esac

function usage()
{
    cat << EOF
Make Linux kernel livepatch.
Usage: $0 -s src-root -b debuginfo-root [-d] [0|1|2]
            -s    kernel/OOT source root, step0 only
            -b    kernel/OOT debuginfo, step0 only
            -d    log debug, optional
            0-2   select step0, 1 or 2 work, optional, default to all
                  step0    collect and appaly patch
                  step1    generate livepatch source
                  step2    build livepatch module

Before running, create a working directory and copy .patch in. The directory
name will be the livepatch module name. In the directory, run klpmake.

Now one livepatch module support only one patch file, either target in-tree
or out-of-tree source.

Suggest running klpmake step by step, check and verify each step's result.
EOF
}

# NOTE: the matching regex might be too limited.
function get_one_module()
{
    local line obj_re ymc_re src_re obj ymc
    obj_re='^[^-]+'
    ymc_re='[^[:blank:]=:+]+'
    src_re="\b$2\b"

    while read line; do
        if [[ "$line" =~ ($obj_re)-($ymc_re).*$src_re ]]; then
            obj=${BASH_REMATCH[1]}
            ymc=${BASH_REMATCH[2]}
            if [[ "$ymc" =~ \$\((.*)\) ]]; then
                ymc=$(grep ${BASH_REMATCH[1]}= /boot/config-`uname -r`|\
                                                            cut -d'=' -f2)
            fi
            if [[ -z $ymc ]]; then
                echo "ERROR: source $1/$2 not enabled in Kconfig"
                return 42
            fi

            # If no obj-* line in the same directory Makefile, then error.
            if [[ "$obj" == "obj" ]]; then
                if [[ $ymc == "y" ]]; then
                    mod="vmlinux"
                else
                    mod="$2"
                fi
                return 0
            else
                get_one_module $1 $obj || return
                return
            fi
        fi
    done <<< $(< $1/Makefile)

    echo "ERROR: failed match $1/$2 in $1/Makefile"
    return 42
}

# Find sources' module they belong to and save in caller's 'modules' array.
function get_modules()
{
    local mod bname dir

    # I believe there must be a Makefile under the source root
    # whether it is kernel or OOT module.
    # If not, I think this scenario is for kernel target and the
    # build files were separated to build path by the distributor.
    if [[ ! -e Makefile ]]; then
        cd /lib/modules/$(uname -r)/build
    fi

    for s in $1; do
        bname=$(basename $s)
        bname=${bname/%.?/.o}
        dir=$(dirname $s)
        get_one_module $dir $bname || return
        modules[$mod]+=$s" "
    done
}

# Find all changed funcs in a patch and save in caller's 'sources' array.
# NOTE: the matching regex might be too limited.
function get_funcs()
{
    local line src_oe name_re at2_re func_re end_re src func saved goon f
    src_re='^--- [^[:blank:]/]+/([^[:blank:]]+)'
    # the last captch group is only for syscall
    name_re='(\b[[:alnum:]_]+\b)[[:blank:]]*\(([^,)]+)'
    at2_re="^@@[^@]+@@.*$name_re"
    func_re="^ [[:alpha:]_].*$name_re"
    end_re='^ }[[:blank:]]*$'

    IFS=
    while read line; do
        if [[ "$line" =~ $src_re ]]; then
            src=${BASH_REMATCH[1]}
            # only care .c file
            if [[ $src == *.c ]]; then
                func=
                saved="no"
                goon="true"
            else
                goon="false"
            fi
        fi
        if [[ $goon == "false" ]]; then
            continue
        fi

        if [[ "$line" =~ $at2_re ]] || [[ "$line" =~ $func_re ]]; then
            f=${BASH_REMATCH[1]}
            # replace syscall macro with real name
            if [[ $f == SYSCALL_DEFINE* ]]; then
                syscalls+=($f"."${BASH_REMATCH[2]})
                f=$SYSCALL_PREFIX${BASH_REMATCH[2]}
            fi
            if [[ $func != $f ]]; then
                func=$f
                saved="no"
            fi
        elif [[ "$line" =~ $end_re ]]; then
            # When matched func close brace, the func is over.
            saved="yes"
        elif [[ "$line" =~ ^(-|\+) ]]; then
            if [[ $saved == "no" ]]; then
                sources[$src]+=$func" "
                saved="yes"
            fi
        fi
    done <<< $(< $1)
    unset IFS
}

function expand_syscall()
{
    for s in $1; do
        s=$(basename $s)
        for sysc in ${syscalls[@]}; do
            awk -f $KLPMAKE_DIR/expand_syscall.awk -v prefix=$SYSCALL_PREFIX \
                -v arg1=$SYSCALL_ARG0 -v arg2=$SYSCALL_ARG1 -v arg3=$SYSCALL_ARG2 \
                -v arg4=$SYSCALL_ARG3 -v arg5=$SYSCALL_ARG4 -v arg6=$SYSCALL_ARG5 \
                -v syscall=$sysc $s > $s.new
            mv -f $s.new $s
        done
    done
}

# apply patch, generate $klpconf for klpsrc
function step0()
{
    echo "step0: collect and apply patch, patch info saved in $klpconf"
    local a_patch srcfiles
    local -A modules sources

    a_patch=$(ls *.patch)
    srcfiles=$(grep "+++ " $a_patch|gawk '{print gensub("[^/]*/", "", 1, $2)}')

    cd $srcroot
    patch --backup --suffix=.klpsrc -p1 --fuzz=0 < $cwd/$a_patch
    for s in $srcfiles; do
        mv -f $s $cwd
        mv $s.klpsrc $s
    done
    get_modules "$srcfiles"

    cd $cwd
    rm -f $klpconf
    get_funcs $a_patch
    echo "obj-m" $(basename $cwd) >> $klpconf
    echo "src-root" $srcroot >> $klpconf
    echo "debug-root" $dbgroot >> $klpconf
    for m in ${!modules[@]}; do
        echo -e "\nmodule-name" $m >> $klpconf
        for s in ${modules[$m]}; do
            echo -e "\tsrc-name" $s >> $klpconf
            for f in ${sources[$s]}; do
                echo -e "\t\tfunc-name" $f >> $klpconf
            done
        done
    done

    if ((${#syscalls[@]} > 0)); then
        expand_syscall  "$srcfiles"
    fi
    echo "OK"
}

# generate livepatch source
# NOTE: source tree must in not-patched state before running
function step1()
{
    echo -e "\nstep1: generate livepatch source, KLPSYMs info saved in $klpsyms"
    local sym

    rm -f $klpsyms
    $KLPMAKE_DIR/klpsrc $debug

    # if there is no KLPSYMs in the livepatch, then it is a normal module
    if [[ -e $klpsyms ]]; then
        sym=$(cut -d' ' -f1 $klpsyms|sort|uniq -d)
        if [[ -n "$sym" ]]; then
            echo "ERROR: not support duplicate KLPSYM: "$sym
            exit 7
        fi
    fi

    for s in $(grep src-name $klpconf|cut -d' ' -f2); do
        s=$(basename $s)
        mv $s $s.patched
        mv $s.klp $s
    done
    echo "OK"
}

# make livepatch module
function step2()
{
    echo -e "\nstep2: making livepatch module...\n"
    local ko

    KBUILD_MODPOST_WARN=1 make
    ko=$(ls *.ko)
    mv -f $ko ${ko}.partial

    if [[ -e $klpsyms ]]; then
        $KLPMAKE_DIR/fixklp ${ko}.partial $klpsyms
    else
        mv -f ${ko}.partial $ko
    fi
    strip -g $ko
    echo "SUCCEED"
}

# parse options
args=`getopt -u -o s:b:d -- $@`
set -- $args
while [[ "$1" != "--" ]]; do
    case "$1" in
        -s)  srcroot=$2; shift 2;;
        -b)  dbgroot=$2; shift 2;;
        -d)  debug=-d; shift;;
         *)  usage; exit 1;;
    esac
done
case $# in
    1)  if [[ -z $srcroot ]] || [[ -z $dbgroot ]]; then
            usage
            exit 2
        fi
        step0
        step1
        step2
        ;;
    2)  if (($2 == 0)); then
            if [[ -z $srcroot ]] || [[ -z $dbgroot ]]; then
                usage
                exit 3
            fi
            step0
        elif (($2 == 1)); then
            step1
        elif (($2 == 2)); then
            step2
        else
            usage
            exit 4
        fi
        ;;
    *)
        usage
        exit 5
        ;;
esac
